import requests
import os
import re
import logging
import json
import csv
from enum import Enum
from dotenv import load_dotenv
from tqdm import tqdm

# 加载 .env 文件中的环境变量
load_dotenv()  # 默认加载当前目录的 .env 文件，若文件路径不同可指定：load_dotenv("/path/to/.env")

# 从环境变量中获取代理配置
http_proxy = os.getenv("HTTP_PROXY")
https_proxy = os.getenv("HTTPS_PROXY")
proxies = {}
if http_proxy:
    proxies["http"] = http_proxy
if https_proxy:
    proxies["https"] = https_proxy

def url_builder(publisher, name, version, platformDiff: bool) -> str:
    """
    组装下载地址
    """
    return f'https://marketplace.visualstudio.com/_apis/public/gallery/publishers/{publisher}/vsextensions/{name}/{version}/vspackage'
      
def preDownload(url) -> tuple:
    """
    通过下载地址判断下载文件名字和文件大小
    return 文件名称, 文件大小, 是否带有平台信息（若带有则需要指定平台）
    """
    try:
        # 发送GET请求，允许重定向（部分资源可能需要跳转）
        response_header = requests.get(url, proxies=proxies, allow_redirects=True, stream=True)
        response_header.raise_for_status()  # 检查请求是否成功（状态码200）

        # 解析Content-Disposition获取文件名
        content_disposition = response_header.headers.get("Content-Disposition", "")
        if "filename=" in content_disposition:
            # 提取文件名（处理带引号或不带引号的情况）
            filename = content_disposition.split("filename=")[-1].strip('"')
            pattern = r'(?:filename=)([a-zA-Z0-9.@%"-]+\.vsix)"?;'
            match = re.search(pattern, content_disposition)
            if match:
                filename = match.group(1).strip('"')
            else:
                logging.warning(f"响应头解析失败，使用默认名字。响应头解析内容：{content_disposition}")
        else:
            # 若响应头无文件名，使用URL解析的默认名称
            logging.warning(f"未从响应头获取到文件名，使用默认名称：{filename}")
            raise FileNotFoundError(f'preDownload 执行失败，未能获取到下载文件名')
        file_size = int(response_header.headers.get('content-length', 0))

        if '@' in filename:
            url += f'?targetPlatform={PlatformType.WIN32_X.value}'

        response_header.close()
        return filename, file_size, url
    except requests.exceptions.RequestException as e:
        print(f"下载失败：{e}")

def doDownload(filename, file_size, url, saved_path):
    response = requests.get(url, proxies=proxies, allow_redirects=True, stream=True)
    response.raise_for_status()
    if not os.path.exists(saved_path):
        os.mkdir(saved_path)
    save_path = os.path.join(saved_path, filename)
    # 写入文件
    with open(save_path, "wb") as f, tqdm(
        desc=filename,
        total=file_size,
        unit='B',
        unit_scale=True,
        unit_divisor=1024,
        miniters=1,
        leave=True
    ) as pbar:
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                f.write(chunk)
                pbar.update(len(chunk))

    print(f"下载成功，文件保存至：{save_path}")

class PlatformType(Enum):
    WIN32_A = "win32-arm64"
    WIN32_X = 'win32-x64'
    LINUX_X = 'linux-x64'
    LINUX_A = 'linux-arm64'
    LINUX_AR = 'linux-armhf'
    ALPINE_A = 'alpine-arm64'
    ALPINE_X = 'alpine-x64'
    DARWIN_X = 'darwin-x64'
    DARWIN_A = 'darwin-arm64'
    WEB = 'Web'

def identifier_analysis(identifier: str, platform:str = PlatformType.WIN32_X.value) -> tuple:
    """
    return 最新版本号, displayName, 是否区分平台
    """
    query_url = r'https://marketplace.visualstudio.com/_apis/public/gallery/extensionquery'
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json;api-version=6.0-preview.1'
    }

    request_json = {"filters": [{"criteria": [{ "filterType": 7, "value": identifier }]}],"flags": 103}
    response = requests.post(url=query_url, headers=headers, json=request_json)
    response_json = json.loads(response.text)

    displayName = response_json['results'][0]['extensions'][0]['displayName']
    version_list = response_json['results'][0]['extensions'][0]['versions']

    for version_info in version_list:
        if 'targetPlatform' in  version_info:
            if version_info['targetPlatform'] == platform:
                return version_info['version'], displayName, True
        else:
            return version_info['version'], displayName, False
    return '', '', False

def load_vsix_list(file_path: str) -> list:
    """
    从CSV文件加载VSIX扩展信息
    """
    vsix_list = []
    file_header = []
    with open(file_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            file_header = reader.fieldnames
            cur_indentifier = row[file_header[0]].strip()
            theVersion, displayName, platformDiff = identifier_analysis(cur_indentifier)
            url = url_builder(cur_indentifier.split('.')[0], cur_indentifier.split('.')[1], theVersion, platformDiff)
            filename, file_size, url = preDownload(url)
            # 更新行数据（覆盖原字段或新增字段）
            row['version'] = theVersion
            row['displayName'] = displayName
            row['fileName'] = filename
            row['fileSize'] = file_size
            row['url'] = url  # 若 CSV 原有该列会覆盖，无该列则新增
            vsix_list.append(row)
    
    with open(file_path, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=file_header)
            writer.writeheader()
            writer.writerows(vsix_list)

    return vsix_list

def batch_download(vsix_list: list, save_path: str):
    """
    批量下载VSIX扩展
    """
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    alreadyDownloadedFileList = os.listdir(save_path)

    for vsix_info in vsix_list:
        # identifier = vsix_info['identifier']
        version = vsix_info['version']
        displayName = vsix_info['displayName']
        fileName = vsix_info['fileName']
        fileSize = vsix_info['fileSize']
        url = vsix_info['url']

        if fileName in alreadyDownloadedFileList:
            print(f"已存在扩展 {displayName} 的版本 {version}，跳过下载。")
            continue

        print(f"正在下载 {displayName} 的版本 {version}...")
        doDownload(fileName, fileSize, url, saved_path=save_path)

if __name__ == '__main__':
    save_path = r'vsix_files'
    csv_file_path = r'vsix_list.csv'
    vsix_list = load_vsix_list(csv_file_path)
    batch_download(vsix_list, save_path)